import os
import torch
from config.dataset import CamVidConfig

class CamVidUnetConfig:
    """Unet model configuration for CamVid dataset - Small real dataset for validation"""
    
    # Model configuration
    MODEL_NAME = "Unet"
    ENCODER_NAME = "resnet34"
    ENCODER_WEIGHTS = None  # Do not use pre-trained weights
    
    # Dataset configuration
    DATASET_CONFIG = CamVidConfig()
    NUM_CLASSES = DATASET_CONFIG.NUM_CLASSES
    
    # Training configuration - Optimized for small dataset
    BATCH_SIZE = 2  # Small batch size for small dataset
    NUM_EPOCHS = 10  # Fewer epochs for quick validation
    LEARNING_RATE = 0.001
    WEIGHT_DECAY = 1e-4
    
    # Image configuration - Use original CamVid size
    IMAGE_SIZE = CamVidConfig.IMAGE_SIZE  # (360, 480)
    
    # Path configuration
    CHECKPOINT_DIR = "checkpoints_camvid"
    LOG_DIR = "logs_camvid"
    
    # Device configuration
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    
    @classmethod
    def create_directories(cls):
        """Create necessary directories"""
        os.makedirs(cls.CHECKPOINT_DIR, exist_ok=True)
        os.makedirs(cls.LOG_DIR, exist_ok=True)


class CamVidSegformerConfig:
    """Segformer model configuration for CamVid dataset - Small real dataset for validation"""
    
    # Model configuration
    MODEL_NAME = "Segformer"
    MODEL_TYPE = "mit_b0"  # Lightweight model
    
    # Dataset configuration
    DATASET_CONFIG = CamVidConfig()
    NUM_CLASSES = DATASET_CONFIG.NUM_CLASSES
    
    # Training configuration - Optimized for small dataset
    BATCH_SIZE = 2  # Small batch size for small dataset
    NUM_EPOCHS = 10  # Fewer epochs for quick validation
    LEARNING_RATE = 0.0001  # Lower learning rate for transformer
    WEIGHT_DECAY = 1e-4
    
    # Image configuration - Use original CamVid size
    IMAGE_SIZE = CamVidConfig.IMAGE_SIZE  # (384, 480)
    
    # Path configuration
    CHECKPOINT_DIR = "checkpoints_camvid_segformer"
    LOG_DIR = "logs_camvid_segformer"
    
    # Device configuration
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    
    @classmethod
    def create_directories(cls):
        """Create necessary directories"""
        os.makedirs(cls.CHECKPOINT_DIR, exist_ok=True)
        os.makedirs(cls.LOG_DIR, exist_ok=True)


# Default configuration for CamVid
CamVidConfig = CamVidUnetConfig